"""
任务存储器 - 内存存储实现
采用策略模式，支持不同的存储策略
"""

import threading
from typing import Dict, List, Optional
from collections import defaultdict
from datetime import datetime

from app.base.base_task_scheduler import BaseTaskStorage
from app.models.task.task_data import Task, TaskStatus, TaskType


class MemoryTaskStorage(BaseTaskStorage):
    """内存任务存储实现"""

    def __init__(self):
        self._tasks: Dict[str, Task] = {}
        self._status_index: Dict[TaskStatus, set] = defaultdict(set)
        self._type_index: Dict[TaskType, set] = defaultdict(set)
        self._lock = threading.RLock()

    async def save_task(self, task: Task) -> None:
        """保存任务"""
        with self._lock:
            self._tasks[task.id] = task
            self._status_index[task.status].add(task.id)
            self._type_index[task.type].add(task.id)

    async def get_task(self, task_id: str) -> Optional[Task]:
        """获取任务"""
        with self._lock:
            return self._tasks.get(task_id)

    async def update_task(self, task: Task) -> None:
        """更新任务"""
        with self._lock:
            old_task = self._tasks.get(task.id)
            if old_task:
                # 更新状态索引
                if old_task.status != task.status:
                    self._status_index[old_task.status].discard(task.id)
                    self._status_index[task.status].add(task.id)

                # 更新类型索引
                if old_task.type != task.type:
                    self._type_index[old_task.type].discard(task.id)
                    self._type_index[task.type].add(task.id)

            self._tasks[task.id] = task

    async def delete_task(self, task_id: str) -> bool:
        """删除任务"""
        with self._lock:
            task = self._tasks.get(task_id)
            if not task:
                return False

            # 从索引中移除
            self._status_index[task.status].discard(task_id)
            self._type_index[task.type].discard(task_id)

            # 删除任务
            del self._tasks[task_id]
            return True

    async def list_tasks(
        self,
        status: Optional[TaskStatus] = None,
        task_type: Optional[TaskType] = None,
        limit: Optional[int] = None,
    ) -> List[Task]:
        """列出任务"""
        with self._lock:
            task_ids = set(self._tasks.keys())

            # 按状态过滤
            if status is not None:
                task_ids &= self._status_index[status]

            # 按类型过滤
            if task_type is not None:
                task_ids &= self._type_index[task_type]

            # 获取任务并排序（按创建时间倒序）
            tasks = [self._tasks[task_id] for task_id in task_ids]
            tasks.sort(key=lambda t: t.created_at, reverse=True)

            # 限制数量
            if limit is not None:
                tasks = tasks[:limit]

            return tasks

    async def get_tasks_count(self) -> Dict[TaskStatus, int]:
        """获取任务统计"""
        with self._lock:
            return {
                status: len(task_ids) for status, task_ids in self._status_index.items()
            }

    async def get_pending_tasks(
        self, task_type: Optional[TaskType] = None
    ) -> List[Task]:
        """获取待执行任务"""
        return await self.list_tasks(status=TaskStatus.PENDING, task_type=task_type)

    async def get_running_tasks(self) -> List[Task]:
        """获取正在执行的任务"""
        return await self.list_tasks(status=TaskStatus.RUNNING)

    async def cleanup_completed_tasks(self, keep_count: int = 1000) -> int:
        """清理已完成的任务，保留最近的指定数量"""
        with self._lock:
            completed_tasks = []
            for task_id in self._status_index[TaskStatus.COMPLETED]:
                task = self._tasks[task_id]
                completed_tasks.append(task)

            # 按完成时间排序，删除旧任务
            completed_tasks.sort(key=lambda t: t.completed_at or datetime.min)
            tasks_to_delete = (
                completed_tasks[:-keep_count]
                if len(completed_tasks) > keep_count
                else []
            )

            deleted_count = 0
            for task in tasks_to_delete:
                if await self.delete_task(task.id):
                    deleted_count += 1

            return deleted_count
